import UIKit
import CoreML
import Vision
import CoreImage

class DenoiseProcessor: ImageProcessor {
    private let model: anime_noise0_model
    
    init() throws {
        self.model = try anime_noise0_model()
    }
    
    func process(image: UIImage) async throws -> UIImage {
        print("开始降噪处理...")
        
        // 1. 创建输入数组
        let shape: [NSNumber] = [3, 142, 142]
        let multiArray = try MLMultiArray(shape: shape, dataType: .float32)
        
        // 2. 调整图片大小
        let targetSize = CGSize(width: 142, height: 142)
        UIGraphicsBeginImageContextWithOptions(targetSize, false, 1.0)
        image.draw(in: CGRect(origin: .zero, size: targetSize))
        guard let resizedImage = UIGraphicsGetImageFromCurrentImageContext(),
              let cgImage = resizedImage.cgImage else {
            UIGraphicsEndImageContext()
            throw ImageProcessingError.resizeFailed
        }
        UIGraphicsEndImageContext()
        
        // 3. 图片数据预处理
        let width = 142
        let height = 142
        let bytesPerPixel = 4
        let bytesPerRow = bytesPerPixel * width
        
        var rawData = [UInt8](repeating: 0, count: height * bytesPerRow)
        let colorSpace = CGColorSpaceCreateDeviceRGB()
        
        guard let context = CGContext(data: &rawData,
                                    width: width,
                                    height: height,
                                    bitsPerComponent: 8,
                                    bytesPerRow: bytesPerRow,
                                    space: colorSpace,
                                    bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else {
            throw ImageProcessingError.invalidImage
        }
        
        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))
        
        // 4. 转换为模型输入格式
        for y in 0..<height {
            for x in 0..<width {
                let offset = (y * bytesPerRow) + (x * bytesPerPixel)
                let r = Float(rawData[offset]) / 255.0
                let g = Float(rawData[offset + 1]) / 255.0
                let b = Float(rawData[offset + 2]) / 255.0
                
                multiArray[[0, y as NSNumber, x as NSNumber] as [NSNumber]] = r as NSNumber
                multiArray[[1, y as NSNumber, x as NSNumber] as [NSNumber]] = g as NSNumber
                multiArray[[2, y as NSNumber, x as NSNumber] as [NSNumber]] = b as NSNumber
            }
        }
        
        // 5. 使用模型进行预测
        let prediction = try model.prediction(input: multiArray)
        let outputArray = prediction.conv7
        
        // 6. 处理输出数据
        let outputWidth = 128
        let outputHeight = 128
        var pixelData = [UInt8](repeating: 0, count: outputWidth * outputHeight * 4)
        
        for y in 0..<outputHeight {
            for x in 0..<outputWidth {
                let offset = (y * outputWidth + x) * 4
                let r = outputArray[[0, 0, 0, y as NSNumber, x as NSNumber] as [NSNumber]].floatValue
                let g = outputArray[[0, 0, 1, y as NSNumber, x as NSNumber] as [NSNumber]].floatValue
                let b = outputArray[[0, 0, 2, y as NSNumber, x as NSNumber] as [NSNumber]].floatValue
                
                pixelData[offset] = UInt8(min(max(r * 255.0, 0), 255))
                pixelData[offset + 1] = UInt8(min(max(g * 255.0, 0), 255))
                pixelData[offset + 2] = UInt8(min(max(b * 255.0, 0), 255))
                pixelData[offset + 3] = 255
            }
        }
        
        // 7. 创建输出图像
        let outputContext = CGContext(data: &pixelData,
                                    width: outputWidth,
                                    height: outputHeight,
                                    bitsPerComponent: 8,
                                    bytesPerRow: outputWidth * 4,
                                    space: colorSpace,
                                    bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue)
        
        guard let outputContext = outputContext,
              let cgImage = outputContext.makeImage() else {
            throw ImageProcessingError.outputImageCreationFailed
        }
        
        let finalImage = UIImage(cgImage: cgImage)
        print("降噪处理完成")
        return finalImage
    }
} 